import numpy as np
from RL.多臂老虎机.Solver import Solver


"""
epsilon贪婪算法: 用小概率ε去套索未知的可能性, 用大概率1-ε去利用当前最优的知识
__init__
    集成多台老虎机的初始化参数和step(使用那台老虎机)方法
    epsilon: ε参数
    estimates: 期望每次使用老虎机每次都是中奖, 数组里面全部是1
run_one_step
    如果随机数小于ε参数, k=随机选择一台老虎机
    否则, k=选择期望值最大的老虎机, 我们初始化的时候期望值全部是1, 所以第一次返回的一定是1
    r(奖励) = BernoulliBandit.step(K), 当前选择的老虎机的奖励值
    更新本次老虎机的estimates期望值估值
    返回当前操作的老虎机
"""
class EpsilonGreedy(Solver):
    """ epsilon贪婪算法,继承Solver类 """
    def __init__(self, bandit, epsilon=0.01, init_prob=1.0):
        super(EpsilonGreedy, self).__init__(bandit)
        self.epsilon = epsilon
        #初始化拉动所有拉杆的期望奖励估值
        self.estimates = np.array([init_prob] * self.bandit.K)

    # 集成父类的run_one_step重新实现
    def run_one_step(self):
        if np.random.random() < self.epsilon:
            k = np.random.randint(0, self.bandit.K)  # 随机选择一根拉杆
        else:
            k = np.argmax(self.estimates)  # 选择期望奖励估值最大的拉杆
        r = self.bandit.step(k)  # 得到本次动作的奖励
        self.estimates[k] += 1. / (self.counts[k] + 1) * (r - self.estimates[k])
        return k
